{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e9f951c0",
   "metadata": {},
   "source": [
    "# TIR VM 代码生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "585be2b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入必要的TVM库和模块\n",
    "import tvm\n",
    "import tvm.testing\n",
    "from tvm import relax\n",
    "from tvm.ir import assert_structural_equal  # 用于比较IR结构是否相等\n",
    "from tvm.script import relax as R  # Relax脚本支持\n",
    "from tvm.script import tir as T  # TIR脚本支持"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8520d54c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 获取TIR模块的辅助函数\n",
    "def get_tir_mod(mod):\n",
    "    \"\"\"\n",
    "    将Relax IR模块转换为TIR模块\n",
    "    \n",
    "    参数:\n",
    "        mod: 输入的Relax IR模块\n",
    "    \n",
    "    返回值:\n",
    "        通过VM代码生成得到的TIR模块\n",
    "    \"\"\"\n",
    "    builder = relax.ExecBuilder()\n",
    "    # 在编译模式下，通过VM代码生成器将Relax模块转换为TIR模块\n",
    "    return relax.vm_build._vmcodegen(builder, mod, exec_mode=\"compiled\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1043185",
   "metadata": {},
   "source": [
    "## 测试VM编译模式下的基本加法算子代码生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ebb45fc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义原始的Relax IR模块\n",
    "@tvm.script.ir_module\n",
    "class Before:\n",
    "    @R.function(pure=False)\n",
    "    def foo(x: R.Tensor):\n",
    "        R.func_attr({\"global_symbol\": \"foo\"})  # 设置全局符号名称\n",
    "        # 调用外部加法函数\n",
    "        z = R.call_packed(\"test.vm.add\", x, x, sinfo_args=(R.Tensor))\n",
    "        return z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b2b09664",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">__vmtir__foo</span>(ctx_ptr: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, r: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, c: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, f: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">2</span>, <span style=\"color: #BA2121\">&quot;test.vm.add&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">0</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">0</span>))\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">1</span>, <span style=\"color: #BA2121\">&quot;vm.builtin.copy&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">2</span>))\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "after = get_tir_mod(Before)\n",
    "after.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "497436bc",
   "metadata": {},
   "source": [
    "## 测试VM编译模式下调用TIR函数的代码生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5bcd8b82",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Before:\n",
    "    # 定义一个TIR原始函数\n",
    "    @T.prim_func\n",
    "    def shape_func(H: T.Buffer(T.int64(4), \"int64\")):\n",
    "        T.func_attr({\"global_symbol\": \"shape_func\"})\n",
    "        # 生成的计算函数：将H的第一个元素加1\n",
    "        H[T.int64(0)] = H[T.int64(0)] + T.int64(1)\n",
    "\n",
    "    @R.function(pure=False)\n",
    "    def foo(x: R.Tensor([4], \"int64\")):\n",
    "        R.func_attr({\"global_symbol\": \"foo\"})\n",
    "        # 调用上面定义的TIR函数\n",
    "        _ = Before.shape_func(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c798c7f3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">__vmtir__foo</span>(ctx_ptr: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, r: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, c: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, f: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>call_cpacked(<span style=\"color: #BA2121\">&quot;shape_func&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">0</span>))\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">1</span>, <span style=\"color: #BA2121\">&quot;vm.builtin.copy&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">0</span>))\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">shape_func</span>(H: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>),), <span style=\"color: #BA2121\">&quot;int64&quot;</span>)):\n",
       "        H[T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">0</span>)] <span style=\"color: #A2F; font-weight: bold\">=</span> H[T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">0</span>)] <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>)\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "after = get_tir_mod(Before)\n",
    "after.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f81c59b4",
   "metadata": {},
   "source": [
    "## 测试VM编译模式下条件分支（if-else）的代码生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "514f9036",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Before:\n",
    "    @R.function(pure=False)\n",
    "    def ife(cond: R.Tensor((), \"bool\"), x: R.Tensor) -> R.Tensor:\n",
    "        R.func_attr({\"global_symbol\": \"ife\"})\n",
    "        # 根据条件执行不同的操作\n",
    "        if cond:\n",
    "            w = R.call_packed(\"test.vm.add\", x, x, sinfo_args=(R.Tensor))\n",
    "        else:\n",
    "            w = R.call_packed(\"test.vm.mul\", x, x, sinfo_args=(R.Tensor))\n",
    "        return w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "946d9f78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">__vmtir__ife</span>(ctx_ptr: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, r: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, c: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, f: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle):\n",
       "        <span style=\"color: #008000; font-weight: bold\">if</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>call_packed(<span style=\"color: #BA2121\">&quot;vm.builtin.read_if_cond&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">0</span>)):\n",
       "            T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">4</span>, <span style=\"color: #BA2121\">&quot;test.vm.add&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">1</span>))\n",
       "            T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">3</span>, <span style=\"color: #BA2121\">&quot;vm.builtin.copy&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">4</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">else</span>:\n",
       "            T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">5</span>, <span style=\"color: #BA2121\">&quot;test.vm.mul&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">1</span>))\n",
       "            T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">3</span>, <span style=\"color: #BA2121\">&quot;vm.builtin.copy&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">5</span>))\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">2</span>, <span style=\"color: #BA2121\">&quot;vm.builtin.copy&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">3</span>))\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "after = get_tir_mod(Before)\n",
    "after.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75b371d0",
   "metadata": {},
   "source": [
    "## 测试VM编译模式下处理常量的代码生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "40870fce",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(x: R.Tensor):\n",
    "        R.func_attr({\"global_symbol\": \"main\"})\n",
    "        y = R.const([1, 2])  # 创建常量数组\n",
    "        z = (y, R.const([3, 4]), x)  # 创建包含常量和变量的元组\n",
    "        return z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "13062860",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">__vmtir__main</span>(ctx_ptr: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, r: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, c: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, f: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">2</span>, <span style=\"color: #BA2121\">&quot;vm.builtin.make_tuple&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(c, <span style=\"color: #008000\">0</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(c, <span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">0</span>))\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">1</span>, <span style=\"color: #BA2121\">&quot;vm.builtin.copy&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">2</span>))\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "after = get_tir_mod(Before)\n",
    "after.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1403a2f",
   "metadata": {},
   "source": [
    "## 测试VM编译模式下使用常量作为函数调用参数的代码生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0491e871",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Before:\n",
    "    @R.function(pure=False)\n",
    "    def main(x: R.Tensor):\n",
    "        R.func_attr({\"global_symbol\": \"main\"})\n",
    "        y = R.const([1, 2])  # 创建常量数组\n",
    "        # 将常量作为参数传递给函数\n",
    "        z = R.call_packed(\"test.vm.add\", x, y, sinfo_args=(R.Tensor))\n",
    "        return z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e4715a24",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">__vmtir__main</span>(ctx_ptr: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, r: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, c: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle, f: T<span style=\"color: #A2F; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">2</span>, <span style=\"color: #BA2121\">&quot;test.vm.add&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">0</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(c, <span style=\"color: #008000\">0</span>))\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_setitem_call_packed(r, <span style=\"color: #008000\">1</span>, <span style=\"color: #BA2121\">&quot;vm.builtin.copy&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>anylist_getitem(r, <span style=\"color: #008000\">2</span>))\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "after = get_tir_mod(Before)\n",
    "after.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9954a4ea",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
